# -*- coding: utf-8 -*-
# @Time : 2023/3/16 14:28
# @Author : Jclian91
# @File : model.py
# @Place : Minghang, Shanghai
import math
import torch
import torch.nn as nn

from params import NUM_WORDS, EMBEDDING_SIZE
from gensim.models import word2vec
import numpy as np
import einops.layers.torch as elt
from encoder import Attention,Attention1

#氨基酸
amino_acids = [str(u) for u in range(0,256)];

def generateWW(dic, word2vec_model):
    dic=list(dic)
    embeddings_matrix = np.zeros((len(dic) + 1, word2vec_model.vector_size),dtype=np.float32)  # 多加了一行，用于padding 补0
    word_index = {'': 0}

    for i in range(len(dic)):
        word = dic[i]
        word_index[word] = i + 1
        embeddings_matrix[i + 1] = word2vec_model.wv[word]
    return word_index, embeddings_matrix

word2vec_model = word2vec.Word2Vec.load('heart_rate_w2c_model.bin')
word_index, embeddings_matrix=generateWW(amino_acids,word2vec_model)

# https://pytorch.org/tutorials/beginner/transformer_tutorial.html
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, vocab_size=10000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(vocab_size, d_model)
        position = torch.arange(0, vocab_size, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float()
            * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer("pe", pe)

    def forward(self, x):
        x = x + self.pe[:, : x.size(1), :]
        return self.dropout(x)


# Text classifier based on a pytorch TransformerEncoder.
class TextClassifier(nn.Module):
    def __init__(
            self,
            nhead=10,
            dim_feedforward=2048,
            num_layers=6,
            dropout=0.1,
            activation="relu",
            classifier_dropout=0.1
    ):
        super().__init__()

        vocab_size = NUM_WORDS + 2
        d_model = EMBEDDING_SIZE
        # vocab_size, d_model = embeddings.size()
        assert d_model % nhead == 0, "nheads must divide evenly into d_model"

        # Embedding layer definition
        self.emb = torch.nn.Embedding.from_pretrained(torch.from_numpy(embeddings_matrix),freeze=False,padding_idx=0)
        #self.emb = nn.Embedding.from_pretrained(pretrained_vector, freeze=False, padding_idx=0)

        self.pos_encoder = PositionalEncoding(
            d_model=d_model,
            dropout=dropout,
            vocab_size=vocab_size
        )

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout
        )
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_layers
        )
        self.classifier = nn.Linear(360, 2)
        self.softmax = nn.LogSoftmax(dim=1)
        self.d_model = d_model
        self.Conv2d = torch.nn.Conv2d(1, 10, kernel_size=[1, 3])
        self.MaxPool2d = torch.nn.MaxPool2d(kernel_size=[1, 5])
        self.Attention = Attention1()
        self.Norm = torch.nn.BatchNorm2d(num_features=10)
        self.relu = torch.nn.ReLU()

    def forward(self, x,y): #x:心率波形数据,y:其他生命指标
        with torch.no_grad():
            x = self.emb(x) * math.sqrt(self.d_model)
        x = self.pos_encoder(x)
        x = self.transformer_encoder(x)
        x = x.mean(dim=1)
        #CNN处理其他特征
        y=y.unsqueeze(1)
        y = y.unsqueeze(0)
        y=elt.Rearrange("b l d k -> l b d k")(y)
        CNR = self.Conv2d(y)

        # CNR=self.relu(CNR)
        # CNR=self.Norm(CNR)
        CMR = self.MaxPool2d(CNR)
        ATT = self.Attention(CMR)
        softATT = torch.softmax(ATT, dim=2)
        RESA = ATT * softATT  # 一个句子的所有词的输出
        mean_RESA = torch.mean(RESA, dim=1)
        mean_RESA = torch.mean(mean_RESA, dim=1)

        final_REs = torch.cat([x, mean_RESA], dim=1)

        x = self.classifier(final_REs)

        return x
